﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Collections;
using System.Windows.Forms;

namespace SnapLib
{
   public static class ProcessExtension
   {
      public static byte[] ReadProcessBytes(this Process p, IntPtr address, int length)
      {
         byte[] buffer = new byte[length];
         int count;
         Win32.ReadProcessMemory(p.Handle, address, buffer, length, out count);
         if (count < length)
            Array.Resize(ref buffer, count);
         return buffer;
      }

      public static int WriteProcessBytes(this Process p, IntPtr address, byte[] buffer, int length)
      {
         int count;
         Win32.WriteProcessMemory(p.Handle, address, buffer, length, out count);
         return count;
      }

      public static IEnumerable<HeapList> Heaps(this Process p)
      {
         IntPtr[] heaps = new IntPtr[250000];
         IntPtr hHeapSnapshot = Toolhelp32.CreateToolhelp32Snapshot(Toolhelp32.SnapshotFlags.All, p.Id);
         Toolhelp32.HEAPLIST32 HeapList = new Toolhelp32.HEAPLIST32();
         HeapList.dwSize = (uint)Marshal.SizeOf(HeapList);
         if (Toolhelp32.Heap32ListFirst(hHeapSnapshot, ref HeapList))
         {
            do
            {
               yield return new HeapList(hHeapSnapshot, HeapList);
            } while (Toolhelp32.Heap32ListNext(hHeapSnapshot, ref HeapList));
         }
         else
         {
            int err = Marshal.GetLastWin32Error();
            MessageBox.Show(err.ToString());
         }
         Win32.CloseHandle(hHeapSnapshot);
      }

   }

   public class HeapList : IEnumerable<Toolhelp32.HEAPENTRY32>
   {
      public Toolhelp32.HEAPLIST32 info { get; private set; }
      public IntPtr hSnapshot { get; private set; }

      public HeapList(IntPtr snapshot, Toolhelp32.HEAPLIST32 info)
      {
         this.info = info;
         this.hSnapshot = snapshot;
      }

      private IEnumerable<Toolhelp32.HEAPENTRY32> Blocks()
      {
         Toolhelp32.HEAPENTRY32 HeapEntry = new Toolhelp32.HEAPENTRY32();
         HeapEntry.dwSize = (uint)Marshal.SizeOf(HeapEntry);
         if (Toolhelp32.Heap32First(ref HeapEntry, info.th32ProcessID, info.th32HeapID))
         {
            do
            {
               yield return HeapEntry;
            } while (Toolhelp32.Heap32Next(out HeapEntry));
         }
         else
         {
            int err = Marshal.GetLastWin32Error();
            MessageBox.Show(err.ToString());
         }

      }

      #region IEnumerable<HEAPENTRY32> Members

      IEnumerator<Toolhelp32.HEAPENTRY32> IEnumerable<Toolhelp32.HEAPENTRY32>.GetEnumerator()
      {
         foreach (Toolhelp32.HEAPENTRY32 entry in Blocks())
            yield return entry;
      }

      IEnumerator System.Collections.IEnumerable.GetEnumerator()
      {
         foreach (Toolhelp32.HEAPENTRY32 entry in Blocks())
            yield return entry;
      }

      #endregion
   }

   public static class Toolhelp32
   {
      public const uint TH32CS_SNAPHEAPLIST = 0x00000001;
      public const uint TH32CS_SNAPPROCESS = 0x00000002;
      public const uint TH32CS_SNAPTHREAD = 0x00000004;
      public const uint TH32CS_SNAPMODULE = 0x00000008;
      public const uint TH32CS_SNAPMODULE32 = 0x00000010;
      public const uint TH32CS_SNAPALL = (TH32CS_SNAPHEAPLIST |
                                               TH32CS_SNAPPROCESS |
                                               TH32CS_SNAPTHREAD |
                                               TH32CS_SNAPMODULE);
      public const uint TH32CS_INHERIT = 0x80000000;

      public const uint HF32_DEFAULT = 1;
      public const uint HF32_SHARED = 2;

      public const uint LF32_FIXED = 0x00000001;
      public const uint LF32_FREE = 0x00000002;
      public const uint LF32_MOVEABLE = 0x00000004;

      public struct HEAPLIST32
      {
         public uint dwSize;
         public uint th32ProcessID;
         public uint th32HeapID;
         public uint dwFlags;
      }

      public struct HEAPENTRY32
      {
         public uint dwSize;
         public IntPtr hHandle;
         public uint dwAddress;
         public uint dwBlockSize;
         public uint dwFlags;
         public uint dwLockCount;
         public uint dwResvd;
         public uint th32ProcessID;
         public uint th32HeapID;
      }

      [Flags]
      public enum SnapshotFlags : uint
      {
         HeapList = 0x00000001,
         Process = 0x00000002,
         Thread = 0x00000004,
         Module = 0x00000008,
         Module32 = 0x00000010,
         Inherit = 0x80000000,
         All = 0x0000001F
      }

      public struct PROCESSENTRY32W
      {
         public uint dwSize;
         public uint cntUsage;
         public uint th32ProcessID;
         public UIntPtr th32DefaultHeapID;
         public uint th32ModuleID;
         public uint cntThreads;
         public uint th32ParentProcessID;
         public int pcPriClassBase;
         public uint dwFlags;
         public string szExeFile;
      }

      public struct PROCESSENTRY32
      {
         public uint dwSize;
         public uint cntUsage;
         public uint th32ProcessID;
         public uint th32DefaultHeapID;
         public uint th32ModuleID;
         public uint cntThreads;
         public uint th32ParentProcessID;
         public int pcPriClassBase;
         public uint dwFlags;
         public string szExeFile;
      }

      public struct THREADENTRY32
      {
         public uint dwSize;
         public uint cntUsage;
         public uint th32ThreadID;
         public uint th32OwnerProcessID;
         public int tpBasePri;
         public int tpDeltaPri;
         public uint dwFlags;
      }

      public struct MODULEENTRY32W
      {
         public uint dwSize;
         public uint th32ModuleID;
         public uint th32ProcessID;
         public uint GlblcntUsage;
         public uint ProccntUsage;
         public IntPtr modBaseAddr;
         public uint modBaseSize;
         public IntPtr hModule;
         public string szModule;
         public string szExePath;
      }

      public struct MODULEENTRY32
      {
         public uint dwSize;
         public uint th32ModuleID;
         public uint th32ProcessID;
         public uint GlblcntUsage;
         public uint ProccntUsage;
         public IntPtr modBaseAddr;
         public uint modBaseSize;
         public IntPtr hModule;
         public string szModule;
         public string szExePath;
      }

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern IntPtr CreateToolhelp32Snapshot(SnapshotFlags dwFlags, int th32ProcessID);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Heap32ListFirst(IntPtr hSnapshot, ref HEAPLIST32 lphl);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Heap32ListNext(IntPtr hSnapshot, ref HEAPLIST32 lphl);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Heap32First(ref HEAPENTRY32 lphe, uint th32ProcessID, uint th32HeapID);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Heap32Next(out HEAPENTRY32 lphe);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Toolhelp32ReadProcessMemory(uint th32ProcessID,
         IntPtr lpBaseAddress, IntPtr lpBuffer, uint cbRead, IntPtr lpNumberOfBytesRead);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Process32FirstW(IntPtr hSnapshot, ref PROCESSENTRY32W lppe);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Process32NextW(IntPtr hSnapshot, ref PROCESSENTRY32W lppe);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Process32First(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Process32Next(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Thread32First(IntPtr hSnapshot, ref THREADENTRY32 lpte);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Thread32Next(IntPtr hSnapshot, ref THREADENTRY32 lpte);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Module32FirstW(IntPtr hSnapshot, ref MODULEENTRY32W lpme);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Module32NextW(IntPtr hSnapshot, ref MODULEENTRY32W lpme);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Module32First(IntPtr hSnapshot, ref MODULEENTRY32W lpme);

      [DllImport("kernel32.dll", SetLastError = true)]
      public static extern bool Module32Next(IntPtr hSnapshot, ref MODULEENTRY32W lpme);
   }

}
